package com.zhang.req.window;

import com.zhang.second.utils.ClickEvent;
import com.zhang.second.utils.ClickSource;
import org.apache.flink.api.common.eventtime.SerializableTimestampAssigner;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.windowing.ProcessWindowFunction;
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.triggers.Trigger;
import org.apache.flink.streaming.api.windowing.triggers.TriggerResult;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.util.Collector;

import java.sql.Timestamp;

/**
 * @title: 10s的滚动窗口但是第一条数据来的下个整数秒每隔一秒输出窗口计算结果
 * @author: zhang
 * @date: 2022/2/3 13:16
 */
public class TriggerWindow {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        env
                .addSource(new ClickSource())
                .assignTimestampsAndWatermarks(
                        WatermarkStrategy.<ClickEvent>forMonotonousTimestamps()
                                .withTimestampAssigner(new SerializableTimestampAssigner<ClickEvent>() {
                                    @Override
                                    public long extractTimestamp(ClickEvent element, long recordTimestamp) {
                                        return element.ts;
                                    }
                                })
                )
                .keyBy(r -> 1)
                .window(TumblingEventTimeWindows.of(Time.seconds(10)))
                .trigger(new MyTrigger())
                .aggregate(new CountAgg(), new WindowResult())
                .print();

        env.execute();
    }

    public static class MyTrigger extends Trigger<ClickEvent, TimeWindow> {
        //窗口中每来一条数据就触发调用一次
        @Override
        public TriggerResult onElement(ClickEvent element, long timestamp, TimeWindow window, TriggerContext ctx) throws Exception {
            ValueState<Boolean> isFirstEvent = ctx.getPartitionedState(
                    new ValueStateDescriptor<Boolean>(
                            "is-first-event",
                            Types.BOOLEAN
                    )
            );
            if (isFirstEvent.value() == null) {
                //属于窗口第一个事件到来
                long nextSecond = element.ts + 1000L - element.ts % 1000L;
                isFirstEvent.update(true);
                ctx.registerEventTimeTimer(nextSecond);
            }
            //窗口什么都不做
            return TriggerResult.CONTINUE;
        }

        //处理时间戳定时器
        //当处理时间到达 'time' 触发onProcessingTime的调用
        @Override
        public TriggerResult onProcessingTime(long time, TimeWindow window, TriggerContext ctx) throws Exception {
            return null;
        }

        //事件时间戳定时器
        //当事件时间到达 'time' 触发onEventTime的调用
        @Override
        public TriggerResult onEventTime(long time, TimeWindow window, TriggerContext ctx) throws Exception {
            if (time < window.getEnd()) {
                if (time + 1000L < window.getEnd()) {
                    ctx.registerEventTimeTimer(time + 1000L);
                }
                //触发后面aggregate的窗口计算
                return TriggerResult.FIRE;
            }
            return TriggerResult.CONTINUE;
        }

        //窗口闭合时候调用
        @Override
        public void clear(TimeWindow window, TriggerContext ctx) throws Exception {
            ValueState<Boolean> isFirstEvent = ctx.getPartitionedState(
                    new ValueStateDescriptor<Boolean>(
                            "is-first-event",
                            Types.BOOLEAN
                    )
            );
            isFirstEvent.clear();
        }
    }

    public static class CountAgg implements AggregateFunction<ClickEvent, Long, Long> {
        @Override
        public Long createAccumulator() {
            return 0L;
        }

        @Override
        public Long add(ClickEvent value, Long accumulator) {
            return accumulator + 1L;
        }

        @Override
        public Long getResult(Long accumulator) {
            return accumulator;
        }

        @Override
        public Long merge(Long a, Long b) {
            return null;
        }
    }

    public static class WindowResult extends ProcessWindowFunction<Long, String, Integer, TimeWindow> {
        @Override
        public void process(Integer integer, ProcessWindowFunction<Long, String, Integer, TimeWindow>.Context context, Iterable<Long> elements, Collector<String> out) throws Exception {
            out.collect("窗口：" + new Timestamp(context.window().getStart()) + "~" + "" +
                    new Timestamp(context.window().getEnd()) + "中pv是：" + elements.iterator().next());
        }
    }
}
